{
    "cells": [
        {
            "cell_type": "markdown",
            "id": "5bf1de44-4047-46cf-a04c-dbf910d9e179",
            "metadata": {},
            "source": [
                "# Ensemble Query Engine Guide\n",
                "\n",
                "Oftentimes when building a RAG applications there are many retreival parameters/strategies to decide from (from chunk size to vector vs. keyword vs. hybrid search, for instance).\n",
                "\n",
                "Thought: what if we could try a bunch of strategies at once, and have any AI/reranker/LLM prune the results?\n",
                "\n",
                "This achieves two purposes:\n",
                "- Better (albeit more costly) retrieved results by pooling results from multiple strategies, assuming the reranker is good\n",
                "- A way to benchmark different retrieval strategies against each other (w.r.t reranker)\n",
                "\n",
                "This guide showcases this over the Great Gatsby. We do ensemble retrieval over different chunk sizes and also different indices.\n",
                "\n",
                "**NOTE**: A closely related guide is our [Ensemble Retrievers Guide](https://gpt-index.readthedocs.io/en/stable/examples/retrievers/ensemble_retrieval.html) - make sure to check it out! "
            ]
        },
        {
            "cell_type": "markdown",
            "id": "6e73fead-ec2c-4346-bd08-e183c13c7e29",
            "metadata": {},
            "source": [
                "## Setup"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 18,
            "id": "a2d59778-4cda-47b5-8cd0-b80fee91d1e4",
            "metadata": {
                "tags": []
            },
            "outputs": [],
            "source": [
                "# NOTE: This is ONLY necessary in jupyter notebook.\n",
                "# Details: Jupyter runs an event-loop behind the scenes.\n",
                "#          This results in nested event-loops when we start an event-loop to make async queries.\n",
                "#          This is normally not allowed, we use nest_asyncio to allow it for convenience.\n",
                "import nest_asyncio\n",
                "\n",
                "nest_asyncio.apply()"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 2,
            "id": "c628448c-573c-4eeb-a7e1-707fe8cc575c",
            "metadata": {
                "tags": []
            },
            "outputs": [
                {
                    "name": "stdout",
                    "output_type": "stream",
                    "text": [
                        "Note: NumExpr detected 12 cores but \"NUMEXPR_MAX_THREADS\" not set, so enforcing safe limit of 8.\n",
                        "NumExpr defaulting to 8 threads.\n"
                    ]
                }
            ],
            "source": [
                "import logging\n",
                "import sys\n",
                "\n",
                "logging.basicConfig(stream=sys.stdout, level=logging.INFO)\n",
                "logging.getLogger().handlers = []\n",
                "logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))\n",
                "\n",
                "from llama_index import (\n",
                "    VectorStoreIndex,\n",
                "    SummaryIndex,\n",
                "    SimpleDirectoryReader,\n",
                "    ServiceContext,\n",
                "    StorageContext,\n",
                "    SimpleKeywordTableIndex,\n",
                ")\n",
                "from llama_index.response.notebook_utils import display_response\n",
                "from llama_index.llms import OpenAI"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "787174ed-10ce-47d7-82fd-9ca7f891eea7",
            "metadata": {},
            "source": [
                "## Load Data\n",
                "\n",
                "We first show how to convert a Document into a set of Nodes, and insert into a DocumentStore."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 3,
            "id": "dd62366b-8a24-40a7-8c47-5859851149fe",
            "metadata": {
                "tags": []
            },
            "outputs": [],
            "source": [
                "# try loading great gatsby\n",
                "\n",
                "documents = SimpleDirectoryReader(\n",
                "    input_files=[\"../../../examples/gatsby/gatsby_full.txt\"]\n",
                ").load_data()"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 4,
            "id": "7081194a-ede7-478e-bff2-23e89e23ef16",
            "metadata": {
                "tags": []
            },
            "outputs": [
                {
                    "name": "stdout",
                    "output_type": "stream",
                    "text": [
                        "Chunk Size: 128\n",
                        "Chunk Size: 256\n",
                        "Chunk Size: 512\n",
                        "Chunk Size: 1024\n"
                    ]
                }
            ],
            "source": [
                "# initialize service context (set chunk size)\n",
                "llm = OpenAI(model=\"gpt-4\")\n",
                "chunk_sizes = [128, 256, 512, 1024]\n",
                "service_contexts = []\n",
                "nodes_list = []\n",
                "vector_indices = []\n",
                "query_engines = []\n",
                "for chunk_size in chunk_sizes:\n",
                "    print(f\"Chunk Size: {chunk_size}\")\n",
                "    service_context = ServiceContext.from_defaults(chunk_size=chunk_size, llm=llm)\n",
                "    service_contexts.append(service_context)\n",
                "    nodes = service_context.node_parser.get_nodes_from_documents(documents)\n",
                "\n",
                "    # add chunk size to nodes to track later\n",
                "    for node in nodes:\n",
                "        node.metadata[\"chunk_size\"] = chunk_size\n",
                "        node.excluded_embed_metadata_keys = [\"chunk_size\"]\n",
                "        node.excluded_llm_metadata_keys = [\"chunk_size\"]\n",
                "\n",
                "    nodes_list.append(nodes)\n",
                "\n",
                "    # build vector index\n",
                "    vector_index = VectorStoreIndex(nodes)\n",
                "    vector_indices.append(vector_index)\n",
                "\n",
                "    # query engines\n",
                "    query_engines.append(vector_index.as_query_engine())"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 5,
            "id": "fbca69b4-d8d5-4dcb-af33-f9ed4a91ec05",
            "metadata": {
                "tags": []
            },
            "outputs": [],
            "source": [
                "# try ensemble retrieval\n",
                "\n",
                "from llama_index.tools import RetrieverTool\n",
                "\n",
                "retriever_tools = []\n",
                "for chunk_size, vector_index in zip(chunk_sizes, vector_indices):\n",
                "    retriever_tool = RetrieverTool.from_defaults(\n",
                "        retriever=vector_index.as_retriever(),\n",
                "        description=f\"Retrieves relevant context from the Great Gatsby (chunk size {chunk_size})\",\n",
                "    )\n",
                "    retriever_tools.append(retriever_tool)"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 6,
            "id": "5c9eaa6f-8f11-4380-b3c6-79092f17def3",
            "metadata": {
                "tags": []
            },
            "outputs": [],
            "source": [
                "from llama_index.selectors.pydantic_selectors import PydanticMultiSelector\n",
                "from llama_index.retrievers import RouterRetriever\n",
                "\n",
                "\n",
                "retriever = RouterRetriever(\n",
                "    selector=PydanticMultiSelector.from_defaults(llm=llm, max_outputs=4),\n",
                "    retriever_tools=retriever_tools,\n",
                ")"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 7,
            "id": "7c72c61c-d4f7-4159-bb80-1989468ab61c",
            "metadata": {
                "tags": []
            },
            "outputs": [
                {
                    "name": "stdout",
                    "output_type": "stream",
                    "text": [
                        "Selecting retriever 0: This choice retrieves a moderate amount of context from the Great Gatsby, which could provide a balanced amount of detail for describing and summarizing the interactions between Gatsby and Daisy..\n",
                        "Selecting retriever 1: This choice retrieves a larger amount of context from the Great Gatsby, which could provide more detail for describing and summarizing the interactions between Gatsby and Daisy..\n",
                        "Selecting retriever 2: This choice retrieves an even larger amount of context from the Great Gatsby, which could provide a comprehensive summary of the interactions between Gatsby and Daisy..\n",
                        "Selecting retriever 3: This choice retrieves the largest amount of context from the Great Gatsby, which could provide the most detailed and comprehensive summary of the interactions between Gatsby and Daisy..\n",
                        "message='OpenAI API response' path=https://api.openai.com/v1/embeddings processing_ms=40 request_id=d269f8a582ac9a70cdb6f587a34d5877 response_code=200\n",
                        "message='OpenAI API response' path=https://api.openai.com/v1/embeddings processing_ms=57 request_id=29679c9e6d594d1f96eb077a4049c6fa response_code=200\n",
                        "message='OpenAI API response' path=https://api.openai.com/v1/embeddings processing_ms=36 request_id=3b3139e7ed9480ff7e7791cc860b6bcd response_code=200\n",
                        "message='OpenAI API response' path=https://api.openai.com/v1/embeddings processing_ms=183 request_id=c0faed36112cb2554de278af9fd58f66 response_code=200\n"
                    ]
                }
            ],
            "source": [
                "nodes = await retriever.aretrieve(\n",
                "    \"Describe and summarize the interactions between Gatsby and Daisy\"\n",
                ")"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 8,
            "id": "590ed8bc-83ad-4851-9ec6-bfbbdf3ff38d",
            "metadata": {
                "scrolled": true,
                "tags": []
            },
            "outputs": [
                {
                    "name": "stdout",
                    "output_type": "stream",
                    "text": [
                        "128\n",
                        "the beach that morning. Finally we came to Gatsby’s own\n",
                        "apartment, a bedroom and a bath, and an Adam’s study, where we sat\n",
                        "down and drank a glass of some Chartreuse he took from a cupboard in\n",
                        "the wall.\n",
                        "\n",
                        "He hadn’t once ceased looking at Daisy, and I think he revalued\n",
                        "everything in his house according to the measure of response it drew\n",
                        "from her well-loved eyes. Sometimes too, he stared around at his\n",
                        "possessions in a dazed\n",
                        "128\n",
                        "turn out as he had\n",
                        "imagined. He had intended, probably, to take what he could and go—but\n",
                        "now he found that he had committed himself to the following of a\n",
                        "grail. He knew that Daisy was extraordinary, but he didn’t realize\n",
                        "just how extraordinary a “nice” girl could be. She vanished into her\n",
                        "rich house, into her rich, full life, leaving Gatsby—nothing. He felt\n",
                        "married to her, that was all.\n",
                        "\n",
                        "When they met again, two days later, it\n",
                        "256\n",
                        "the\n",
                        "direction. In this heat every extra gesture was an affront to the\n",
                        "common store of life.\n",
                        "\n",
                        "The room, shadowed well with awnings, was dark and cool. Daisy and\n",
                        "Jordan lay upon an enormous couch, like silver idols weighing down\n",
                        "their own white dresses against the singing breeze of the fans.\n",
                        "\n",
                        "“We can’t move,” they said together.\n",
                        "\n",
                        "Jordan’s fingers, powdered white over their tan, rested for a moment\n",
                        "in mine.\n",
                        "\n",
                        "“And Mr. Thomas Buchanan, the athlete?” I inquired.\n",
                        "\n",
                        "Simultaneously I heard his voice, gruff, muffled, husky, at the hall\n",
                        "telephone.\n",
                        "\n",
                        "Gatsby stood in the centre of the crimson carpet and gazed around with\n",
                        "fascinated eyes. Daisy watched him and laughed, her sweet, exciting\n",
                        "laugh; a tiny gust of powder rose from her bosom into the air.\n",
                        "\n",
                        "“The rumour is,” whispered Jordan, “that\n",
                        "256\n",
                        "In the meantime, In between time—”\n",
                        "\n",
                        "As I went over to say goodbye I saw that the expression of\n",
                        "bewilderment had come back into Gatsby’s face, as though a faint doubt\n",
                        "had occurred to him as to the quality of his present happiness. Almost\n",
                        "five years! There must have been moments even that afternoon when\n",
                        "Daisy tumbled short of his dreams—not through her own fault, but\n",
                        "because of the colossal vitality of his illusion. It had gone beyond\n",
                        "her, beyond everything. He had thrown himself into it with a creative\n",
                        "passion, adding to it all the time, decking it out with every bright\n",
                        "feather that drifted his way. No amount of fire or freshness can\n",
                        "challenge what a man can store up in his ghostly heart.\n",
                        "\n",
                        "As I watched him he adjusted himself a little, visibly. His hand took\n",
                        "hold of hers, and as she said something low in his ear he turned\n",
                        "toward her with a rush of emotion. I think that voice held him most,\n",
                        "with its fluctuating,\n",
                        "512\n",
                        "go downstairs,” interrupted Gatsby. He flipped a switch. The\n",
                        "grey windows disappeared as the house glowed full of light.\n",
                        "\n",
                        "In the music-room Gatsby turned on a solitary lamp beside the piano.\n",
                        "He lit Daisy’s cigarette from a trembling match, and sat down with her\n",
                        "on a couch far across the room, where there was no light save what the\n",
                        "gleaming floor bounced in from the hall.\n",
                        "\n",
                        "When Klipspringer had played “The Love Nest” he turned around on the\n",
                        "bench and searched unhappily for Gatsby in the gloom.\n",
                        "\n",
                        "“I’m all out of practice, you see. I told you I couldn’t play. I’m all\n",
                        "out of prac—”\n",
                        "\n",
                        "“Don’t talk so much, old sport,” commanded Gatsby. “Play!”\n",
                        "\n",
                        " “In the morning, In the evening, Ain’t we got fun—”\n",
                        "\n",
                        "Outside the wind was loud and there was a faint flow of thunder along\n",
                        "the Sound. All the lights were going on in West Egg now; the electric\n",
                        "trains, men-carrying, were plunging home through the rain from New\n",
                        "York. It was the hour of a profound human change, and excitement was\n",
                        "generating on the air.\n",
                        "\n",
                        " “One thing’s sure and nothing’s surer The rich get richer and the\n",
                        " poor get—children. In the meantime, In between time—”\n",
                        "\n",
                        "As I went over to say goodbye I saw that the expression of\n",
                        "bewilderment had come back into Gatsby’s face, as though a faint doubt\n",
                        "had occurred to him as to the quality of his present happiness. Almost\n",
                        "five years! There must have been moments even that afternoon when\n",
                        "Daisy tumbled short of his dreams—not through her own fault, but\n",
                        "because of the colossal vitality of his illusion. It had gone beyond\n",
                        "her, beyond everything. He had thrown himself into it with a creative\n",
                        "passion, adding to it all the time, decking it out\n",
                        "512\n",
                        "world complete\n",
                        "in itself, with its own standards and its own great figures, second to\n",
                        "nothing because it had no consciousness of being so, and now I was\n",
                        "looking at it again, through Daisy’s eyes. It is invariably saddening\n",
                        "to look through new eyes at things upon which you have expended your\n",
                        "own powers of adjustment.\n",
                        "\n",
                        "They arrived at twilight, and, as we strolled out among the sparkling\n",
                        "hundreds, Daisy’s voice was playing murmurous tricks in her throat.\n",
                        "\n",
                        "“These things excite me so,” she whispered. “If you want to kiss me\n",
                        "any time during the evening, Nick, just let me know and I’ll be glad\n",
                        "to arrange it for you. Just mention my name. Or present a green card.\n",
                        "I’m giving out green—”\n",
                        "\n",
                        "“Look around,” suggested Gatsby.\n",
                        "\n",
                        "“I’m looking around. I’m having a marvellous—”\n",
                        "\n",
                        "“You must see the faces of many people you’ve heard about.”\n",
                        "\n",
                        "Tom’s arrogant eyes roamed the crowd.\n",
                        "\n",
                        "“We don’t go around very much,” he said; “in fact, I was just thinking\n",
                        "I don’t know a soul here.”\n",
                        "\n",
                        "“Perhaps you know that lady.” Gatsby indicated a gorgeous, scarcely\n",
                        "human orchid of a woman who sat in state under a white-plum tree. Tom\n",
                        "and Daisy stared, with that peculiarly unreal feeling that accompanies\n",
                        "the recognition of a hitherto ghostly celebrity of the movies.\n",
                        "\n",
                        "“She’s lovely,” said Daisy.\n",
                        "\n",
                        "“The man bending over her is her director.”\n",
                        "\n",
                        "He took them ceremoniously from group to group:\n",
                        "\n",
                        "“Mrs. Buchanan … and Mr. Buchanan—” After an instant’s hesitation he\n",
                        "added: “the polo player.”\n",
                        "\n",
                        "“Oh no,” objected Tom quickly,\n",
                        "1024\n",
                        "The\n",
                        "grey windows disappeared as the house glowed full of light.\n",
                        "\n",
                        "In the music-room Gatsby turned on a solitary lamp beside the piano.\n",
                        "He lit Daisy’s cigarette from a trembling match, and sat down with her\n",
                        "on a couch far across the room, where there was no light save what the\n",
                        "gleaming floor bounced in from the hall.\n",
                        "\n",
                        "When Klipspringer had played “The Love Nest” he turned around on the\n",
                        "bench and searched unhappily for Gatsby in the gloom.\n",
                        "\n",
                        "“I’m all out of practice, you see. I told you I couldn’t play. I’m all\n",
                        "out of prac—”\n",
                        "\n",
                        "“Don’t talk so much, old sport,” commanded Gatsby. “Play!”\n",
                        "\n",
                        " “In the morning, In the evening, Ain’t we got fun—”\n",
                        "\n",
                        "Outside the wind was loud and there was a faint flow of thunder along\n",
                        "the Sound. All the lights were going on in West Egg now; the electric\n",
                        "trains, men-carrying, were plunging home through the rain from New\n",
                        "York. It was the hour of a profound human change, and excitement was\n",
                        "generating on the air.\n",
                        "\n",
                        " “One thing’s sure and nothing’s surer The rich get richer and the\n",
                        " poor get—children. In the meantime, In between time—”\n",
                        "\n",
                        "As I went over to say goodbye I saw that the expression of\n",
                        "bewilderment had come back into Gatsby’s face, as though a faint doubt\n",
                        "had occurred to him as to the quality of his present happiness. Almost\n",
                        "five years! There must have been moments even that afternoon when\n",
                        "Daisy tumbled short of his dreams—not through her own fault, but\n",
                        "because of the colossal vitality of his illusion. It had gone beyond\n",
                        "her, beyond everything. He had thrown himself into it with a creative\n",
                        "passion, adding to it all the time, decking it out with every bright\n",
                        "feather that drifted his way. No amount of fire or freshness can\n",
                        "challenge what a man can store up in his ghostly heart.\n",
                        "\n",
                        "As I watched him he adjusted himself a little, visibly. His hand took\n",
                        "hold of hers, and as she said something low in his ear he turned\n",
                        "toward her with a rush of emotion. I think that voice held him most,\n",
                        "with its fluctuating, feverish warmth, because it couldn’t be\n",
                        "over-dreamed—that voice was a deathless song.\n",
                        "\n",
                        "They had forgotten me, but Daisy glanced up and held out her hand;\n",
                        "Gatsby didn’t know me now at all. I looked once more at them and they\n",
                        "looked back at me, remotely, possessed by intense life. Then I went\n",
                        "out of the room and down the marble steps into the rain, leaving them\n",
                        "there together.\n",
                        "\n",
                        "\n",
                        "                                  VI\n",
                        "\n",
                        "About this time an ambitious young reporter from New York arrived one\n",
                        "morning at Gatsby’s door and asked him if he had anything to say.\n",
                        "\n",
                        "“Anything to say about what?” inquired Gatsby politely.\n",
                        "\n",
                        "“Why—any statement to give out.”\n",
                        "\n",
                        "It transpired after a confused five minutes that the man had heard\n",
                        "Gatsby’s name around his office in a connection which he either\n",
                        "wouldn’t reveal or didn’t fully understand. This was his day off and\n",
                        "with laudable initiative he had hurried out “to see.”\n",
                        "\n",
                        "It was a random shot, and yet the reporter’s instinct was right.\n",
                        "Gatsby’s notoriety, spread about by the hundreds who had accepted his\n",
                        "hospitality and so become authorities upon his past, had increased all\n",
                        "summer until he fell just short of being news. Contemporary legends\n",
                        "such as the “underground pipeline to Canada” attached themselves to\n",
                        "him, and there was one persistent story that he didn’t live in a house\n",
                        "at all, but in a boat that looked like a house and was\n",
                        "1024\n",
                        "Daisy insistently. Gatsby’s eyes\n",
                        "floated toward her. “Ah,” she cried, “you look so cool.”\n",
                        "\n",
                        "Their eyes met, and they stared together at each other, alone in\n",
                        "space. With an effort she glanced down at the table.\n",
                        "\n",
                        "“You always look so cool,” she repeated.\n",
                        "\n",
                        "She had told him that she loved him, and Tom Buchanan saw. He was\n",
                        "astounded. His mouth opened a little, and he looked at Gatsby, and\n",
                        "then back at Daisy as if he had just recognized her as someone he knew\n",
                        "a long time ago.\n",
                        "\n",
                        "“You resemble the advertisement of the man,” she went on innocently.\n",
                        "“You know the advertisement of the man—”\n",
                        "\n",
                        "“All right,” broke in Tom quickly, “I’m perfectly willing to go to\n",
                        "town. Come on—we’re all going to town.”\n",
                        "\n",
                        "He got up, his eyes still flashing between Gatsby and his wife. No one\n",
                        "moved.\n",
                        "\n",
                        "“Come on!” His temper cracked a little. “What’s the matter, anyhow?\n",
                        "If we’re going to town, let’s start.”\n",
                        "\n",
                        "His hand, trembling with his effort at self-control, bore to his lips\n",
                        "the last of his glass of ale. Daisy’s voice got us to our feet and out\n",
                        "on to the blazing gravel drive.\n",
                        "\n",
                        "“Are we just going to go?” she objected. “Like this? Aren’t we going\n",
                        "to let anyone smoke a cigarette first?”\n",
                        "\n",
                        "“Everybody smoked all through lunch.”\n",
                        "\n",
                        "“Oh, let’s have fun,” she begged him. “It’s too hot to fuss.”\n",
                        "\n",
                        "He didn’t answer.\n",
                        "\n",
                        "“Have it your own way,” she said. “Come on, Jordan.”\n",
                        "\n",
                        "They went upstairs to get ready while we three men stood there\n",
                        "shuffling the hot pebbles with our feet. A silver curve of the moon\n",
                        "hovered already in the western sky. Gatsby started to speak, changed\n",
                        "his mind, but not before Tom wheeled and faced him expectantly.\n",
                        "\n",
                        "“Have you got your stables here?” asked Gatsby with an effort.\n",
                        "\n",
                        "“About a quarter of a mile down the road.”\n",
                        "\n",
                        "“Oh.”\n",
                        "\n",
                        "A pause.\n",
                        "\n",
                        "“I don’t see the idea of going to town,” broke out Tom savagely.\n",
                        "“Women get these notions in their heads—”\n",
                        "\n",
                        "“Shall we take anything to drink?” called Daisy from an upper window.\n",
                        "\n",
                        "“I’ll get some whisky,” answered Tom. He went inside.\n",
                        "\n",
                        "Gatsby turned to me rigidly:\n",
                        "\n",
                        "“I can’t say anything in his house, old sport.”\n",
                        "\n",
                        "“She’s got an indiscreet voice,” I remarked. “It’s full of—” I\n",
                        "hesitated.\n",
                        "\n",
                        "“Her voice is full of money,” he said suddenly.\n",
                        "\n",
                        "That was it. I’d never understood before. It was full of money—that\n",
                        "was the inexhaustible charm that rose and fell in it, the jingle of\n",
                        "it, the cymbals’ song of it … High in a white palace the king’s\n",
                        "daughter, the golden girl …\n",
                        "\n",
                        "Tom came out of the house wrapping a quart bottle in a towel, followed\n",
                        "by Daisy and Jordan wearing small tight hats of metallic cloth and\n",
                        "carrying light capes over their arms.\n",
                        "\n",
                        "“Shall we all go in my car?” suggested Gatsby. He felt the hot, green\n",
                        "leather of the seat. “I ought to have left it in the shade.”\n",
                        "\n",
                        "“Is it standard shift?” demanded Tom.\n",
                        "\n",
                        "“Yes.”\n",
                        "\n",
                        "“Well, you take my coupé and let me drive your car to town.”\n",
                        "\n",
                        "The suggestion was distasteful to Gatsby.\n",
                        "\n",
                        "“I\n"
                    ]
                }
            ],
            "source": [
                "for node in nodes:\n",
                "    print(node.node.metadata[\"chunk_size\"])\n",
                "    print(node.node.get_text())"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 9,
            "id": "1f26c527-17d2-4d4e-a6ee-8ea878ef8742",
            "metadata": {
                "tags": []
            },
            "outputs": [],
            "source": [
                "# define reranker\n",
                "from llama_index.indices.postprocessor import (\n",
                "    LLMRerank,\n",
                "    SentenceTransformerRerank,\n",
                "    CohereRerank,\n",
                ")\n",
                "\n",
                "# reranker = LLMRerank()\n",
                "# reranker = SentenceTransformerRerank(top_n=10)\n",
                "reranker = CohereRerank(top_n=10)"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 10,
            "id": "828589ef-d062-40dc-8a4b-245190769445",
            "metadata": {
                "tags": []
            },
            "outputs": [],
            "source": [
                "# define RetrieverQueryEngine\n",
                "from llama_index.query_engine import RetrieverQueryEngine\n",
                "\n",
                "query_engine = RetrieverQueryEngine(retriever, node_postprocessors=[reranker])"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 11,
            "id": "53e3c341-e66d-4950-88d5-6411699d064b",
            "metadata": {
                "tags": []
            },
            "outputs": [
                {
                    "name": "stdout",
                    "output_type": "stream",
                    "text": [
                        "Selecting retriever 0: This choice provides a moderate chunk size that could contain relevant interactions between Gatsby and Daisy without being too overwhelming..\n",
                        "Selecting retriever 1: This choice provides a larger chunk size that could contain more detailed interactions between Gatsby and Daisy..\n",
                        "Selecting retriever 2: This choice provides an even larger chunk size that could contain extensive interactions between Gatsby and Daisy, providing a more comprehensive summary..\n",
                        "Selecting retriever 3: This choice provides the largest chunk size that could contain the most detailed and comprehensive interactions between Gatsby and Daisy, but it might also include a lot of irrelevant information..\n"
                    ]
                }
            ],
            "source": [
                "response = query_engine.query(\n",
                "    \"Describe and summarize the interactions between Gatsby and Daisy\"\n",
                ")"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "9aa680dd-03a0-4a76-b456-c4ef0136fdc2",
            "metadata": {
                "tags": []
            },
            "outputs": [],
            "source": [
                "display_response(\n",
                "    response, show_source=True, source_length=500, show_source_metadata=True\n",
                ")"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 13,
            "id": "0a7a8303-be94-45c5-8bc5-13ec8c7f1694",
            "metadata": {
                "tags": []
            },
            "outputs": [],
            "source": [
                "# compute the average precision for each chunk size based on positioning in combined ranking\n",
                "from collections import defaultdict\n",
                "import pandas as pd\n",
                "\n",
                "\n",
                "def mrr_all(metadata_values, metadata_key, source_nodes):\n",
                "    # source nodes is a ranked list\n",
                "    # go through each value, find out positioning in source_nodes\n",
                "    value_to_mrr_dict = {}\n",
                "    for metadata_value in metadata_values:\n",
                "        mrr = 0\n",
                "        for idx, source_node in enumerate(source_nodes):\n",
                "            if source_node.node.metadata[metadata_key] == metadata_value:\n",
                "                mrr = 1 / (idx + 1)\n",
                "                break\n",
                "            else:\n",
                "                continue\n",
                "\n",
                "        # normalize AP, set in dict\n",
                "        value_to_mrr_dict[metadata_value] = mrr\n",
                "\n",
                "    df = pd.DataFrame(value_to_mrr_dict, index=[\"MRR\"])\n",
                "    df.style.set_caption(\"Mean Reciprocal Rank\")\n",
                "    return df"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 14,
            "id": "adebbb82-764e-4b45-933e-84bf4ad64d40",
            "metadata": {
                "tags": []
            },
            "outputs": [
                {
                    "name": "stdout",
                    "output_type": "stream",
                    "text": [
                        "Mean Reciprocal Rank for each Chunk Size\n"
                    ]
                },
                {
                    "data": {
                        "text/html": [
                            "<div>\n",
                            "<style scoped>\n",
                            "    .dataframe tbody tr th:only-of-type {\n",
                            "        vertical-align: middle;\n",
                            "    }\n",
                            "\n",
                            "    .dataframe tbody tr th {\n",
                            "        vertical-align: top;\n",
                            "    }\n",
                            "\n",
                            "    .dataframe thead th {\n",
                            "        text-align: right;\n",
                            "    }\n",
                            "</style>\n",
                            "<table border=\"1\" class=\"dataframe\">\n",
                            "  <thead>\n",
                            "    <tr style=\"text-align: right;\">\n",
                            "      <th></th>\n",
                            "      <th>128</th>\n",
                            "      <th>256</th>\n",
                            "      <th>512</th>\n",
                            "      <th>1024</th>\n",
                            "    </tr>\n",
                            "  </thead>\n",
                            "  <tbody>\n",
                            "    <tr>\n",
                            "      <th>MRR</th>\n",
                            "      <td>0.2</td>\n",
                            "      <td>0.166667</td>\n",
                            "      <td>0.5</td>\n",
                            "      <td>1.0</td>\n",
                            "    </tr>\n",
                            "  </tbody>\n",
                            "</table>\n",
                            "</div>"
                        ],
                        "text/plain": [
                            "     128       256   512   1024\n",
                            "MRR   0.2  0.166667   0.5   1.0"
                        ]
                    },
                    "execution_count": 14,
                    "metadata": {},
                    "output_type": "execute_result"
                }
            ],
            "source": [
                "# Compute the Mean Reciprocal Rank for each chunk size (higher is better)\n",
                "# we can see that chunk size of 256 has the highest ranked results.\n",
                "print(\"Mean Reciprocal Rank for each Chunk Size\")\n",
                "mrr_all(chunk_sizes, \"chunk_size\", response.source_nodes)"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "b27a2f3c-55ce-4fa6-a15a-be539723a967",
            "metadata": {},
            "source": [
                "## Compare Against Baseline\n",
                "\n",
                "Compare against a baseline of chunk size 1024 (k=2)"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 15,
            "id": "a4d66b14-4f38-4b61-809c-f603d7e09ef9",
            "metadata": {
                "tags": []
            },
            "outputs": [],
            "source": [
                "query_engine_1024 = query_engines[-1]"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 16,
            "id": "43f3e441-f372-4df2-ae21-71fa7968e606",
            "metadata": {
                "tags": []
            },
            "outputs": [],
            "source": [
                "response_1024 = query_engine_1024.query(\n",
                "    \"Describe and summarize the interactions between Gatsby and Daisy\"\n",
                ")"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "4b161e69-da17-4e4e-b8c0-b9c846ce723f",
            "metadata": {
                "tags": []
            },
            "outputs": [],
            "source": [
                "display_response(response_1024, show_source=True, source_length=500)"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "3e7cea30-b036-4a8b-b16d-823ab6b860be",
            "metadata": {},
            "outputs": [],
            "source": []
        }
    ],
    "metadata": {
        "kernelspec": {
            "display_name": "llama_index_v2",
            "language": "python",
            "name": "llama_index_v2"
        },
        "language_info": {
            "codemirror_mode": {
                "name": "ipython",
                "version": 3
            },
            "file_extension": ".py",
            "mimetype": "text/x-python",
            "name": "python",
            "nbconvert_exporter": "python",
            "pygments_lexer": "ipython3",
            "version": "3.10.10"
        }
    },
    "nbformat": 4,
    "nbformat_minor": 5
}